package org.example.yuaiagent.chatmemory;

import com.alibaba.fastjson2.JSON;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import org.example.yuaiagent.model.entity.MyChatMessage;
import org.example.yuaiagent.service.ChatMessageService;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;
import java.util.UUID;
@Component
public class MysqlChatMemory implements ChatMemory {

    @Autowired
    private ChatMessageService chatMessageService;

    @Override
    public void add(String conversationId, List<Message> messages) {

        List<MyChatMessage> mYMessages =
                messages.stream()
                        .map(message -> {
                            Map<String, Object> metadata = message.getMetadata();
                            metadata.putIfAbsent("timestamp", System.currentTimeMillis());
                            MyChatMessage myChatMessage = new MyChatMessage(UUID.randomUUID().toString(), conversationId, JSON.toJSONString(message));
                            myChatMessage.setMessageType(message.getMessageType());
                            myChatMessage.setTextContent(message.getText());
                            myChatMessage.setMetadata(JSON.toJSONString(metadata));
                            return myChatMessage;
    }).toList();
        chatMessageService.saveBatch(mYMessages);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        //按照时间取最新的N条
        LambdaQueryWrapper<MyChatMessage> wrapper = new LambdaQueryWrapper<>();
        wrapper.eq(MyChatMessage::getConversationId, conversationId).orderByDesc(MyChatMessage::getCreateTime).last("limit " + lastN);
        List<Message> list = chatMessageService.list(wrapper).stream().map(message -> JSON.parseObject(message.getMessage(), Message.class)).toList();
        List<Message> ret = list.stream().map(message -> {
            Message msg=null;
            if (message.getMessageType().getValue().equals("assistant")) {
                //转化为 assistant
                msg = new AssistantMessage(message.getText());
            }
            if (message.getMessageType().getValue().equals("user")) {
                msg = new UserMessage(message.getText());
            }
            if (message.getMessageType().getValue().equals("system")) {
                msg = new SystemMessage(message.getText());
            }

            if (msg != null) {
                BeanUtils.copyProperties(message, msg);
                // 手动复制 metadata 内容
                Map<String, Object> metadata = message.getMetadata();
                if (metadata != null) {
                    msg.getMetadata().putAll(metadata);  // 将 metadata 内容复制到 msg 中
                }
            }
            return msg;
        }).toList();
        return ret;
    }

    @Override
    public void clear(String conversationId) {
        chatMessageService.remove(new LambdaQueryWrapper<MyChatMessage>().eq(MyChatMessage::getConversationId, conversationId));
    }

    public List<Message> getContextByPage(String conversationId, int pageNum, int pageSize) {
        LambdaQueryWrapper<MyChatMessage> wrapper = new LambdaQueryWrapper<MyChatMessage>()
                .eq(MyChatMessage::getConversationId, conversationId)
                .orderByDesc(MyChatMessage::getCreateTime)
                .last("LIMIT " + (pageNum-1)*pageSize + "," + pageSize); // 加入分页参数
        return chatMessageService.list(wrapper).stream().map(msg -> {
            return JSON.parseObject(msg.getMessage(), Message.class);
        }).toList();
    }
}